import os
import os.path as osp

import cv2
import numpy as np


def get_img_list(path):
    img_list = []
    for root, dirs, files in os.walk(path):
        for f in files:
            if f.endswith('.jpg') and f.find('_mask') == -1:
                img_list.append(osp.join(root, f))

    assert len(img_list) > 0, "没有图片！"
    return img_list


if __name__ == "__main__":
    img_path = r"data\shushu\train"

    img_list = get_img_list(img_path)

    for img_path in img_list:
        mask_path = img_path.replace('.jpg', '_mask.png').replace('train', 'train_heatmap')
        assert osp.exists(mask_path), "path is not exist"
        img = cv2.imread(img_path)
        mask_inp = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
        # print("mask unique: ", np.unique(mask_inp))

        mask = np.zeros_like(img)
        mask[..., 1] = mask_inp

        img_show = img.copy()
        img_show[mask!=0] = np.array(img_show[mask!=0] * 0.5 + mask[mask!=0] * 0.5, np.uint8)
        # img_show[:, :, 0] = np.array(img_show[:, :, 0] + mask * 0.5, np.uint8)

        img_show = cv2.cvtColor(img_show, cv2.COLOR_BGR2RGB)

        cv2.imshow('show', img_show)
        cv2.waitKey()

        _=1

